#ifndef NN_MATCHING_H
#define NN_MATCHING_H

#include <map>
#include <vector>
#include <opencv2/opencv.hpp>
#include "utils.h"
#include <numeric>
#include <omp.h>
#include "op_cosinesimilarity.h"

using namespace std;

vector<vector<float>> iou_dist(const vector<BBox> &dets, const vector<BBox> &trks);

// save features of the track in GPU
class FeatureBundle {
public:
    FeatureBundle() : store(budget, vector<float>(feat_dim, 0)), full(false), next(0) {}

    void clear() {
        next = 0;
        full = false;
    }

    bool empty() const {
        return next == 0 && !full;
    }

    void add(vector<float> feat) {
        if (next == budget) {
            full = true;
            next = 0;
        }
        store[next++] = feat;
    }

    vector<vector<float>> get() const {
        return vector<vector<float>>(store.begin(), store.begin()+next);
    }

private:
    static const int64_t budget = 100, feat_dim = 512;

    vector<vector<float>> store;

    bool full;
    int64_t next;
};

template<typename TrackData>
class FeatureMetric {
public:
    explicit FeatureMetric(vector<TrackData> &data) : data(data) {
    }

    vector<vector<float>> distance(vector<vector<float>>& features, const vector<int> &targets) {
        // static bool inited = false;
        // if (!inited) {
        //     if (ATLAS_OK != op.Init()) {
        //         ATLAS_LOG_ERROR("init op failed");
        //     }
        //     inited = true;
        // }
        struct timeval begin;
        struct timeval end;
        gettimeofday(&begin, NULL);

        auto dist = vector<vector<float>>(targets.size(), vector<float>(features.size(), 0));
        if (features.size()) {
            #pragma omp parallel for schedule(dynamic)
            for (size_t i = 0; i < targets.size(); ++i) {
                dist[i] = nn_cosine_distance(data[targets[i]].feats.get(), features);
            }
        }

        // Result ret = op.RunCosineSimilarityOp(v, s,o);
        // if (ret != SUCCESS) {
        //     ATLAS_LOG_ERROR("op run failed");
        // }
        


        gettimeofday(&end, NULL);
        float time_cost = 1000 * (end.tv_sec - begin.tv_sec) + (end.tv_usec - begin.tv_usec) / 1000.000;
        INFO_LOG("NN distance time: %fms", time_cost);
        return dist;
    }

    void update(vector<vector<float>>& feature, const vector<int> &targets) {
        for (size_t i = 0; i < targets.size(); ++i) {
            data[targets[i]].feats.add(feature[i]);
        }
    }

private:
    std::vector<TrackData> &data;
    // OpCosineSimilarity op;

    vector<float> nn_cosine_distance(const vector<vector<float>>& x, const vector<vector<float>>& y) {
        vector<float> res(y.size(), 0);
#pragma omp parallel for schedule(dynamic)
        for (size_t i = 0; i < y.size(); ++i) {
            double min_dist = 1;
            // double minop = 1;
#pragma omp parallel for reduction(min:min_dist)
            for (size_t j = 0; j < x.size(); ++j) {
                // vector<void*> ov;
                // vector<void*> iv;
                // vector<size_t> sv;
                // iv.push_back((void*)x[j].data());
                // iv.push_back((void*)y[i].data());
                // sv.push_back(x[j].size());
                // sv.push_back(y[i].size());

                // op.RunCosineSimilarityOp(iv, sv, ov);
                // minop = std::min(1.0-((float*)(ov[0]))[0], minop);
                min_dist = std::min(1.0-std::inner_product(std::begin(x[j]), std::end(x[j]), std::begin(y[i]), 0.0), min_dist);
            }
            res[i] = (float)min_dist;
        }
        return res;
    }
};

#endif //NN_MATCHING_H
